package datastructure.find;

import java.util.LinkedList;
import java.util.Queue;

/**
 * @author MaoLin Wang
 * @description TODO
 * @date 2020/8/26 4:42 下午
 */
public class AVLTree<T extends Comparable<? super T>> {
    private AvlNode<T> root;
    private static final int DEFAULT_HEIGHT=1;

    public AVLTree(AvlNode<T> root) {
        this.root = root;
    }
    public void remove(T data){
        if (data==null){
            return;
        }
       root= remove(data,root);
    }

    public void insert(T data){
        if (data==null){
            return;
        }
        root=insert(data,root);
    }

    private AvlNode<T> insert(T data, AvlNode<T> root) {
        if (root==null){
            return new AvlNode<>(data);
        }
        int result=data.compareTo(root.data);
        if (result<0){
            root.lchild=insert(data,root.lchild);
        }else if (result>0){
            root.rchild=insert(data,root.rchild);
        }else {
            ;
        }
        return balance(root);
    }

    private AvlNode<T> remove(T data, AvlNode<T> root) {
        if (root==null){
            return null;
        }
        int result=root.data.compareTo(data);
        if (result<0){
            root.rchild=remove(data,root.rchild);
        }else if (result>0){
            root.lchild=remove(data,root.lchild);
        }else {
            if (root.lchild!=null && root.rchild!=null){
                root.data=findMin(root).data;
                root.rchild=remove(root.data,root.rchild);
            }else {
                root=root.lchild==null?root.rchild:root.lchild;
            }
        }
        return balance(root);
    }

    public Queue<T>iterator(){
        if (root!=null){
            Queue<T>queue=new LinkedList<>();
            return iterator(root,queue);
        }
        return null;
    }

    private Queue<T> iterator(AvlNode<T> root,Queue<T>queue) {
        if (root!=null){
            iterator(root.lchild,queue);
            queue.add(root.data);
            iterator(root.rchild,queue);
        }
        return queue;
    }

    private AvlNode<T> balance(AvlNode<T> root) {
        if (root==null){
            return null;
        }
        if (height(root.lchild)-height(root.rchild)>DEFAULT_HEIGHT){
            if (height(root.lchild.lchild)>=height(root.lchild.rchild)){
                //左子树的左孩子添加
                root=rotateRight(root);
            }else {
                root=doubleRotateLeftRight(root);
            }
        }else if (height(root.rchild)-height(root.lchild)>DEFAULT_HEIGHT){
            if (height(root.rchild.rchild)>=height(root.rchild.lchild)){
                root=rotateLeft(root);
            }else {
                root =doubleRotateRightLeft(root);
            }
        }
        root.height=Math.max(height(root.lchild),height(root.rchild))+1;
        return root;
    }

    private AvlNode<T> doubleRotateRightLeft(AvlNode<T> root) {
        root.rchild=rotateRight(root.rchild);
        return rotateLeft(root);
    }

    private AvlNode<T> doubleRotateLeftRight(AvlNode<T> root) {
        root.lchild=rotateLeft(root.lchild);
        return rotateRight(root);
    }

    private AvlNode<T> rotateLeft(AvlNode<T> root) {
        AvlNode<T>right=root.rchild;
        root.rchild=right.lchild;
        right.lchild=root;
        root.height=1+Math.max(height(root.lchild),height(root.rchild));
        right.height=1+Math.max(height(right.rchild),root.height);
        return right;
    }

    private AvlNode<T> rotateRight(AvlNode<T> root) {
        AvlNode<T>left=root.lchild;
        root.lchild=left.rchild;
        left.rchild=root;
        root.height=Math.max(height(root.lchild),height(root.rchild))+1;
        left.height=Math.max(height(root.lchild),root.height)+1;
        return left;
    }

    private int height(AvlNode<T> root) {
        return root==null?-1:root.height;
    }

    public AvlNode<T> findMin(){
        if (root==null){
            return null;
        }
        return findMin(root);
    }

    private AvlNode<T> findMin(AvlNode<T> root) {
        if (root==null){
            return null;
        }
        if (root.lchild==null){
            return root.rchild;
        }
        while (root.lchild!=null) {
            root = root.lchild;
        }
        return root;
    }

    public static void main(String[] args) {
        AVLTree<Integer> avlTree = new AVLTree<>(new AvlNode<>(1));
        avlTree.insert(2);
        avlTree.insert(3);
        avlTree.insert(4);
        avlTree.insert(5);
        avlTree.insert(6);
        avlTree.insert(7);
        avlTree.insert(15);
        avlTree.insert(16);
        avlTree.insert(17);

        avlTree.insert(14);
        //中序遍历打印
        avlTree.printTree();
        System.out.println("==========删除节点15=========");

        avlTree.remove(15);
        Queue<Integer> queue = avlTree.iterator();
        while (!queue.isEmpty()){
            System.out.println(queue.poll());
        }
        Integer max = avlTree.findMin().data;
        System.out.println("最大元素:"+max);
        System.out.println("节点17是否存在"+avlTree.contains(17));
    }

    private boolean contains(T data) {
        if (data==null){
            return false;
        }
        if (root!=null){
            return contains(data,root);
        }
        return false;
    }

    private boolean contains(T data, AvlNode<T> root) {
        if (root==null){
            return false;
        }
        int result=data.compareTo(root.data);
        if (result<0){
            return contains(data,root.lchild);
        }else if (result>0){
            return contains(data,root.rchild);
        }else {
            return true;
        }
    }

    private void printTree() {
        if (root==null){
            System.out.println("空");
        }else {
            printTree(root);
        }
    }

    private void printTree(AvlNode<T> root) {
        if (root!=null){
            printTree(root.lchild);
            System.out.println(root.data);
            printTree(root.rchild);
        }

    }
}

class AvlNode<T>{
    T data;
    AvlNode<T> lchild;
    AvlNode<T> rchild;
    int height;

    public AvlNode(T data) {
        this(data,null,null);
    }

    public AvlNode(T data, AvlNode<T> lchild, AvlNode<T> rchild) {
        this.data = data;
        this.lchild = lchild;
        this.rchild = rchild;
    }
}